{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install datasets accelerate transformers torch>=2.1.2 evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "zVvslsfMIrIh"
   },
   "outputs": [],
   "source": [
    "# 根据你使用的模型和GPU资源情况，调整以下关键参数\n",
    "squad_v2 = True\n",
    "model_checkpoint = \"distilbert-base-uncased\"\n",
    "batch_size = 16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "IreSlFmlIrIm"
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 270,
     "referenced_widgets": [
      "69caab03d6264fef9fc5649bffff5e20",
      "3f74532faa86412293d90d3952f38c4a",
      "50615aa59c7247c4804ca5cbc7945bd7",
      "fe962391292a413ca55dc932c4279fa7",
      "299f4b4c07654e53a25f8192bd1d7bbd",
      "ad04ed1038154081bbb0c1444784dcc2",
      "7c667ad22b5740d5a6319f1b1e3a8097",
      "46c2b043c0f84806978784a45a4e203b",
      "80e2943be35f46eeb24c8ab13faa6578",
      "de5956b5008d4fdba807bae57509c393",
      "931db1f7a42f4b46b7ff8c2e1262b994",
      "6c1db72efff5476e842c1386fadbbdba",
      "ccd2f37647c547abb4c719b75a26f2de",
      "d30a66df5c0145e79693e09789d96b81",
      "5fa26fc336274073abbd1d550542ee33",
      "2b34de08115d49d285def9269a53f484",
      "d426be871b424affb455aeb7db5e822e",
      "160bf88485f44f5cb6eaeecba5e0901f",
      "745c0d47d672477b9bb0dae77b926364",
      "d22ab78269cd4ccfbcf70c707057c31b",
      "d298eb19eeff453cba51c2804629d3f4",
      "a7204ade36314c86907c562e0a2158b8",
      "e35d42b2d352498ca3fc8530393786b2",
      "75103f83538d44abada79b51a1cec09e",
      "f6253931d90543e9b5fd0bb2d615f73a",
      "051aa783ff9e47e28d1f9584043815f5",
      "0984b2a14115454bbb009df71c1cf36f",
      "8ab9dfce29854049912178941ef1b289",
      "c9de740e007141958545e269372780a4",
      "cbea68b25d6d4ba09b2ce0f27b1726d5",
      "5781fc45cf8d486cb06ed68853b2c644",
      "d2a92143a08a4951b55bab9bc0a6d0d3",
      "a14c3e40e5254d61ba146f6ec88eae25",
      "c4ffe6f624ce4e978a0d9b864544941a",
      "1aca01c1d8c940dfadd3e7144bb35718",
      "9fbbaae50e6743f2aa19342152398186",
      "fea27ca6c9504fc896181bc1ff5730e5",
      "940d00556cb849b3a689d56e274041c2",
      "5cdf9ed939fb42d4bf77301c80b8afca",
      "94b39ccfef0b4b08bf2fb61bb0a657c1",
      "9a55087c85b74ea08b3e952ac1d73cbe",
      "2361ab124daf47cc885ff61f2899b2af",
      "1a65887eb37747ddb75dc4a40f7285f2",
      "3c946e2260704e6c98593136bd32d921",
      "50d325cdb9844f62a9ecc98e768cb5af",
      "aa781f0cfe454e9da5b53b93e9baabd8",
      "6bb68d3887ef43809eb23feb467f9723",
      "7e29a8b952cf4f4ea42833c8bf55342f",
      "dd5997d01d8947e4b1c211433969b89b",
      "2ace4dc78e2f4f1492a181bcd63304e7",
      "bbee008c2791443d8610371d1f16b62b",
      "31b1c8a2e3334b72b45b083688c1a20c",
      "7fb7c36adc624f7dbbcb4a831c1e4f63",
      "0b7c8f1939074794b3d9221244b1344d",
      "a71908883b064e1fbdddb547a8c41743",
      "2f5223f26c8541fc87e91d2205c39995"
     ]
    },
    "id": "s_AY1ATSIrIq",
    "outputId": "fd0578d1-8895-443d-b56f-5908de9f1b6b"
   },
   "outputs": [],
   "source": [
    "datasets = load_dataset(\"/root/autodl-tmp/datasets/squad_v2\" if squad_v2 else \"squad\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RzfPtOMoIrIu"
   },
   "source": [
    "The `datasets` object itself is [`DatasetDict`](https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasetdict), which contains one key for the training, validation and test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "GWiVUF0jIrIv",
    "outputId": "35e3ea43-f397-4a54-c90c-f2cf8d36873e"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['id', 'title', 'context', 'question', 'answers'],\n",
       "        num_rows: 130319\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['id', 'title', 'context', 'question', 'answers'],\n",
       "        num_rows: 11873\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The maximum length of a feature (question and context)\n",
    "max_length = 384 \n",
    "# The authorized overlap between two part of the context when splitting it is needed.\n",
    "doc_stride = 128 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "pad_on_right = tokenizer.padding_side == \"right\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 原始训练过程，删减部分代码，只留相关结果数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "id": "uNx5pyRlIrJh",
    "outputId": "077e661e-d36c-469b-89b8-7ff7f73541ec"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='4152' max='4152' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [4152/4152 2:23:19, Epoch 3/3]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>1.491100</td>\n",
       "      <td>1.249441</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>1.108800</td>\n",
       "      <td>1.161671</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>0.975700</td>\n",
       "      <td>1.158766</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Checkpoint destination directory models/distilbert-base-uncased-finetuned-squad/checkpoint-500 already exists and is non-empty.Saving will proceed but saved results may be invalid.\n",
      "Checkpoint destination directory models/distilbert-base-uncased-finetuned-squad/checkpoint-1000 already exists and is non-empty.Saving will proceed but saved results may be invalid.\n",
      "Checkpoint destination directory models/distilbert-base-uncased-finetuned-squad/checkpoint-1500 already exists and is non-empty.Saving will proceed but saved results may be invalid.\n",
      "Checkpoint destination directory models/distilbert-base-uncased-finetuned-squad/checkpoint-2000 already exists and is non-empty.Saving will proceed but saved results may be invalid.\n",
      "Checkpoint destination directory models/distilbert-base-uncased-finetuned-squad/checkpoint-2500 already exists and is non-empty.Saving will proceed but saved results may be invalid.\n",
      "Checkpoint destination directory models/distilbert-base-uncased-finetuned-squad/checkpoint-3000 already exists and is non-empty.Saving will proceed but saved results may be invalid.\n",
      "Checkpoint destination directory models/distilbert-base-uncased-finetuned-squad/checkpoint-3500 already exists and is non-empty.Saving will proceed but saved results may be invalid.\n",
      "Checkpoint destination directory models/distilbert-base-uncased-finetuned-squad/checkpoint-4000 already exists and is non-empty.Saving will proceed but saved results may be invalid.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=4152, training_loss=1.3038662743246854, metrics={'train_runtime': 8602.4737, 'train_samples_per_second': 30.872, 'train_steps_per_second': 0.483, 'total_flos': 2.602335381127373e+16, 'train_loss': 1.3038662743246854, 'epoch': 3.0})"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 模型评估"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'exact_match': 74.88174077578051, 'f1': 83.6359321422016}"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "if squad_v2:\n",
    "    formatted_predictions = [{\"id\": k, \"prediction_text\": v, \"no_answer_probability\": 0.0} for k, v in final_predictions.items()]\n",
    "else:\n",
    "    formatted_predictions = [{\"id\": k, \"prediction_text\": v} for k, v in final_predictions.items()]\n",
    "references = [{\"id\": ex[\"id\"], \"answers\": ex[\"answers\"]} for ex in datasets[\"validation\"]]\n",
    "metric.compute(predictions=formatted_predictions, references=references)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Homework：加载本地保存的模型，进行评估和再训练更高的 F1 Score"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. 数据准备"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "# The maximum length of a feature (question and context)\n",
    "max_length = 384 \n",
    "# The authorized overlap between two part of the context when splitting it is needed.\n",
    "doc_stride = 128 \n",
    "\n",
    "model_checkpoint = f\"/root/autodl-tmp/models/distilbert-base-uncased\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_train_features(examples):\n",
    "    # 一些问题的左侧可能有很多空白字符，这对我们没有用，而且会导致上下文的截断失败\n",
    "    # （标记化的问题将占用大量空间）。因此，我们删除左侧的空白字符。\n",
    "    examples[\"question\"] = [q.lstrip() for q in examples[\"question\"]]\n",
    "\n",
    "    pad_on_right = tokenizer.padding_side == \"right\"\n",
    "\n",
    "    # 使用截断和填充对我们的示例进行标记化，但保留溢出部分，使用步幅（stride）。\n",
    "    # 当上下文很长时，这会导致一个示例可能提供多个特征，其中每个特征的上下文都与前一个特征的上下文有一些重叠。\n",
    "    tokenized_examples = tokenizer(\n",
    "        examples[\"question\" if pad_on_right else \"context\"],\n",
    "        examples[\"context\" if pad_on_right else \"question\"],\n",
    "        truncation=\"only_second\" if pad_on_right else \"only_first\",\n",
    "        max_length=max_length,\n",
    "        stride=doc_stride,\n",
    "        return_overflowing_tokens=True,\n",
    "        return_offsets_mapping=True,\n",
    "        padding=\"max_length\",\n",
    "    )\n",
    "\n",
    "    # 由于一个示例可能给我们提供多个特征（如果它具有很长的上下文），我们需要一个从特征到其对应示例的映射。这个键就提供了这个映射关系。\n",
    "    sample_mapping = tokenized_examples.pop(\"overflow_to_sample_mapping\")\n",
    "    # 偏移映射将为我们提供从令牌到原始上下文中的字符位置的映射。这将帮助我们计算开始位置和结束位置。\n",
    "    offset_mapping = tokenized_examples.pop(\"offset_mapping\")\n",
    "\n",
    "    # 让我们为这些示例进行标记！\n",
    "    tokenized_examples[\"start_positions\"] = []\n",
    "    tokenized_examples[\"end_positions\"] = []\n",
    "\n",
    "    for i, offsets in enumerate(offset_mapping):\n",
    "        # 我们将使用 CLS 特殊 token 的索引来标记不可能的答案。\n",
    "        input_ids = tokenized_examples[\"input_ids\"][i]\n",
    "        cls_index = input_ids.index(tokenizer.cls_token_id)\n",
    "\n",
    "        # 获取与该示例对应的序列（以了解上下文和问题是什么）。\n",
    "        sequence_ids = tokenized_examples.sequence_ids(i)\n",
    "\n",
    "        # 一个示例可以提供多个跨度，这是包含此文本跨度的示例的索引。\n",
    "        sample_index = sample_mapping[i]\n",
    "        answers = examples[\"answers\"][sample_index]\n",
    "        # 如果没有给出答案，则将cls_index设置为答案。\n",
    "        if len(answers[\"answer_start\"]) == 0:\n",
    "            tokenized_examples[\"start_positions\"].append(cls_index)\n",
    "            tokenized_examples[\"end_positions\"].append(cls_index)\n",
    "        else:\n",
    "            # 答案在文本中的开始和结束字符索引。\n",
    "            start_char = answers[\"answer_start\"][0]\n",
    "            end_char = start_char + len(answers[\"text\"][0])\n",
    "\n",
    "            # 当前跨度在文本中的开始令牌索引。\n",
    "            token_start_index = 0\n",
    "            while sequence_ids[token_start_index] != (1 if pad_on_right else 0):\n",
    "                token_start_index += 1\n",
    "\n",
    "            # 当前跨度在文本中的结束令牌索引。\n",
    "            token_end_index = len(input_ids) - 1\n",
    "            while sequence_ids[token_end_index] != (1 if pad_on_right else 0):\n",
    "                token_end_index -= 1\n",
    "\n",
    "            # 检测答案是否超出跨度（在这种情况下，该特征的标签将使用CLS索引）。\n",
    "            if not (offsets[token_start_index][0] <= start_char and offsets[token_end_index][1] >= end_char):\n",
    "                tokenized_examples[\"start_positions\"].append(cls_index)\n",
    "                tokenized_examples[\"end_positions\"].append(cls_index)\n",
    "            else:\n",
    "                # 否则，将token_start_index和token_end_index移到答案的两端。\n",
    "                # 注意：如果答案是最后一个单词（边缘情况），我们可以在最后一个偏移之后继续。\n",
    "                while token_start_index < len(offsets) and offsets[token_start_index][0] <= start_char:\n",
    "                    token_start_index += 1\n",
    "                tokenized_examples[\"start_positions\"].append(token_start_index - 1)\n",
    "                while offsets[token_end_index][1] >= end_char:\n",
    "                    token_end_index -= 1\n",
    "                tokenized_examples[\"end_positions\"].append(token_end_index + 1)\n",
    "\n",
    "    return tokenized_examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dd38845033674c8a8c24bfcc749be87b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/11873 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "tokenized_datasets = datasets.map(prepare_train_features,\n",
    "                                  batched=True,\n",
    "                                  remove_columns=datasets[\"train\"].column_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForQuestionAnswering, TrainingArguments, Trainer\n",
    "batch_size=64\n",
    "model_dir = f\"/root/autodl-tmp/output/models/{model_checkpoint}-finetuned-squad\"\n",
    "\n",
    "args = TrainingArguments(\n",
    "    output_dir=model_dir,\n",
    "    evaluation_strategy = \"epoch\",\n",
    "    learning_rate=2e-5,\n",
    "    per_device_train_batch_size=batch_size,\n",
    "    per_device_eval_batch_size=batch_size,\n",
    "    num_train_epochs=3,\n",
    "    weight_decay=0.01,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of DistilBertForQuestionAnswering were not initialized from the model checkpoint at /root/autodl-tmp/models/distilbert-base-uncased and are newly initialized: ['qa_outputs.bias', 'qa_outputs.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "model_checkpoint = f\"/root/autodl-tmp/models/distilbert-base-uncased\"\n",
    "trained_model = AutoModelForQuestionAnswering.from_pretrained(model_checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import default_data_collator\n",
    "\n",
    "data_collator = default_data_collator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    trained_model,\n",
    "    args,\n",
    "    train_dataset=tokenized_datasets[\"train\"],\n",
    "    eval_dataset=tokenized_datasets[\"validation\"],\n",
    "    data_collator=data_collator,\n",
    "    tokenizer=tokenizer,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='6177' max='6177' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [6177/6177 52:23, Epoch 3/3]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>1.423100</td>\n",
       "      <td>1.295846</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>1.158000</td>\n",
       "      <td>1.253980</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>1.022100</td>\n",
       "      <td>1.339991</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=6177, training_loss=1.3404366763564053, metrics={'train_runtime': 3144.1763, 'train_samples_per_second': 125.712, 'train_steps_per_second': 1.965, 'total_flos': 3.873165421863629e+16, 'train_loss': 1.3404366763564053, 'epoch': 3.0})"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_to_save = trainer.save_model(model_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "odict_keys(['loss', 'start_logits', 'end_logits'])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "for batch in trainer.get_eval_dataloader():\n",
    "    break\n",
    "batch = {k: v.to(trainer.args.device) for k, v in batch.items()}\n",
    "with torch.no_grad():\n",
    "    output = trainer.model(**batch)\n",
    "output.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_best_size = 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "start_logits = output.start_logits[0].cpu().numpy()\n",
    "end_logits = output.end_logits[0].cpu().numpy()\n",
    "\n",
    "# 获取最佳的起始和结束位置的索引：\n",
    "start_indexes = np.argsort(start_logits)[-1 : -n_best_size - 1 : -1].tolist()\n",
    "end_indexes = np.argsort(end_logits)[-1 : -n_best_size - 1 : -1].tolist()\n",
    "\n",
    "valid_answers = []\n",
    "\n",
    "# 遍历起始位置和结束位置的索引组合\n",
    "for start_index in start_indexes:\n",
    "    for end_index in end_indexes:\n",
    "        if start_index <= end_index:  # 需要进一步测试以检查答案是否在上下文中\n",
    "            valid_answers.append(\n",
    "                {\n",
    "                    \"score\": start_logits[start_index] + end_logits[end_index],\n",
    "                    \"text\": \"\"  # 我们需要找到一种方法来获取与上下文中答案对应的原始子字符串\n",
    "                }\n",
    "            )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "pad_on_right = tokenizer.padding_side == \"right\"\n",
    "def prepare_validation_features(examples):\n",
    "    # 一些问题的左侧有很多空白，这些空白并不有用且会导致上下文截断失败（分词后的问题会占用很多空间）。\n",
    "    # 因此我们移除这些左侧空白\n",
    "    examples[\"question\"] = [q.lstrip() for q in examples[\"question\"]]\n",
    "\n",
    "    # 使用截断和可能的填充对我们的示例进行分词，但使用步长保留溢出的令牌。这导致一个长上下文的示例可能产生\n",
    "    # 几个特征，每个特征的上下文都会稍微与前一个特征的上下文重叠。\n",
    "    tokenized_examples = tokenizer(\n",
    "        examples[\"question\" if pad_on_right else \"context\"],\n",
    "        examples[\"context\" if pad_on_right else \"question\"],\n",
    "        truncation=\"only_second\" if pad_on_right else \"only_first\",\n",
    "        max_length=max_length,\n",
    "        stride=doc_stride,\n",
    "        return_overflowing_tokens=True,\n",
    "        return_offsets_mapping=True,\n",
    "        padding=\"max_length\",\n",
    "    )\n",
    "\n",
    "    # 由于一个示例在上下文很长时可能会产生几个特征，我们需要一个从特征映射到其对应示例的映射。这个键就是为了这个目的。\n",
    "    sample_mapping = tokenized_examples.pop(\"overflow_to_sample_mapping\")\n",
    "\n",
    "    # 我们保留产生这个特征的示例ID，并且会存储偏移映射。\n",
    "    tokenized_examples[\"example_id\"] = []\n",
    "\n",
    "    for i in range(len(tokenized_examples[\"input_ids\"])):\n",
    "        # 获取与该示例对应的序列（以了解哪些是上下文，哪些是问题）。\n",
    "        sequence_ids = tokenized_examples.sequence_ids(i)\n",
    "        context_index = 1 if pad_on_right else 0\n",
    "\n",
    "        # 一个示例可以产生几个文本段，这里是包含该文本段的示例的索引。\n",
    "        sample_index = sample_mapping[i]\n",
    "        tokenized_examples[\"example_id\"].append(examples[\"id\"][sample_index])\n",
    "\n",
    "        # 将不属于上下文的偏移映射设置为None，以便容易确定一个令牌位置是否属于上下文。\n",
    "        tokenized_examples[\"offset_mapping\"][i] = [\n",
    "            (o if sequence_ids[k] == context_index else None)\n",
    "            for k, o in enumerate(tokenized_examples[\"offset_mapping\"][i])\n",
    "        ]\n",
    "\n",
    "    return tokenized_examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d0ba7513f2514b2799780c899b7cd912",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/11873 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "validation_features = datasets[\"validation\"].map(\n",
    "    prepare_validation_features,\n",
    "    batched=True,\n",
    "    remove_columns=datasets[\"validation\"].column_names\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "raw_predictions = trainer.predict(validation_features)\n",
    "validation_features.set_format(type=validation_features.format[\"type\"], columns=list(validation_features.features.keys()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_answer_length = 30"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'score': 13.722256, 'text': 'France'},\n",
       " {'score': 8.904267, 'text': 'France.'},\n",
       " {'score': 7.4105177,\n",
       "  'text': 'France. They were descended from Norse (\"Norman\" comes from \"Norseman\") raiders and pirates from Denmark, Iceland and Norway'},\n",
       " {'score': 7.1157217, 'text': 'in France'},\n",
       " {'score': 6.9141216,\n",
       "  'text': 'France. They were descended from Norse (\"Norman\" comes from \"Norseman\") raiders and pirates from Denmark'},\n",
       " {'score': 6.3763213, 'text': 'a region in France'},\n",
       " {'score': 6.1270814, 'text': 'Normandy, a region in France'},\n",
       " {'score': 5.9484673,\n",
       "  'text': 'France. They were descended from Norse (\"Norman\" comes from \"Norseman\") raiders and pirates from Denmark, Iceland'},\n",
       " {'score': 5.6872473, 'text': 'region in France'},\n",
       " {'score': 3.9132898,\n",
       "  'text': 'French: Normands; Latin: Normanni) were the people who in the 10th and 11th centuries gave their name to Normandy, a region in France'},\n",
       " {'score': 3.3954978,\n",
       "  'text': '10th and 11th centuries gave their name to Normandy, a region in France'},\n",
       " {'score': 3.1347907, 'text': ', a region in France'},\n",
       " {'score': 2.773508,\n",
       "  'text': 'France. They were descended from Norse (\"Norman\" comes from \"Norseman\") raiders and pirates from Denmark,'},\n",
       " {'score': 2.297733, 'text': 'in France.'},\n",
       " {'score': 1.5583324, 'text': 'a region in France.'},\n",
       " {'score': 1.3090925, 'text': 'Normandy, a region in France.'},\n",
       " {'score': 1.0843282, 'text': 'Denmark, Iceland and Norway'},\n",
       " {'score': 0.8692585, 'text': 'region in France.'},\n",
       " {'score': 0.80398333,\n",
       "  'text': 'in France. They were descended from Norse (\"Norman\" comes from \"Norseman\") raiders and pirates from Denmark, Iceland and Norway'},\n",
       " {'score': 0.5879321, 'text': 'Denmark'}]"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "start_logits = output.start_logits[0].cpu().numpy()\n",
    "end_logits = output.end_logits[0].cpu().numpy()\n",
    "offset_mapping = validation_features[0][\"offset_mapping\"]\n",
    "\n",
    "# 第一个特征来自第一个示例。对于更一般的情况，我们需要将example_id匹配到一个示例索引\n",
    "context = datasets[\"validation\"][0][\"context\"]\n",
    "\n",
    "# 收集最佳开始/结束逻辑的索引：\n",
    "start_indexes = np.argsort(start_logits)[-1 : -n_best_size - 1 : -1].tolist()\n",
    "end_indexes = np.argsort(end_logits)[-1 : -n_best_size - 1 : -1].tolist()\n",
    "valid_answers = []\n",
    "for start_index in start_indexes:\n",
    "    for end_index in end_indexes:\n",
    "        # 不考虑超出范围的答案，原因是索引超出范围或对应于输入ID的部分不在上下文中。\n",
    "        if (\n",
    "            start_index >= len(offset_mapping)\n",
    "            or end_index >= len(offset_mapping)\n",
    "            or offset_mapping[start_index] is None\n",
    "            or offset_mapping[end_index] is None\n",
    "        ):\n",
    "            continue\n",
    "        # 不考虑长度小于0或大于max_answer_length的答案。\n",
    "        if end_index < start_index or end_index - start_index + 1 > max_answer_length:\n",
    "            continue\n",
    "        if start_index <= end_index: # 我们需要细化这个测试，以检查答案是否在上下文中\n",
    "            start_char = offset_mapping[start_index][0]\n",
    "            end_char = offset_mapping[end_index][1]\n",
    "            valid_answers.append(\n",
    "                {\n",
    "                    \"score\": start_logits[start_index] + end_logits[end_index],\n",
    "                    \"text\": context[start_char: end_char]\n",
    "                }\n",
    "            )\n",
    "\n",
    "valid_answers = sorted(valid_answers, key=lambda x: x[\"score\"], reverse=True)[:n_best_size]\n",
    "valid_answers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.auto import tqdm\n",
    "import collections\n",
    "def postprocess_qa_predictions(examples, features, raw_predictions, n_best_size = 20, max_answer_length = 30):\n",
    "    all_start_logits, all_end_logits = raw_predictions\n",
    "    # 构建一个从示例到其对应特征的映射。\n",
    "    example_id_to_index = {k: i for i, k in enumerate(examples[\"id\"])}\n",
    "    features_per_example = collections.defaultdict(list)\n",
    "    for i, feature in enumerate(features):\n",
    "        features_per_example[example_id_to_index[feature[\"example_id\"]]].append(i)\n",
    "\n",
    "    # 我们需要填充的字典。\n",
    "    predictions = collections.OrderedDict()\n",
    "\n",
    "    # 日志记录。\n",
    "    print(f\"正在后处理 {len(examples)} 个示例的预测，这些预测分散在 {len(features)} 个特征中。\")\n",
    "\n",
    "    # 遍历所有示例！\n",
    "    for example_index, example in enumerate(tqdm(examples)):\n",
    "        # 这些是与当前示例关联的特征的索引。\n",
    "        feature_indices = features_per_example[example_index]\n",
    "\n",
    "        min_null_score = None # 仅在squad_v2为True时使用。\n",
    "        valid_answers = []\n",
    "        \n",
    "        context = example[\"context\"]\n",
    "        # 遍历与当前示例关联的所有特征。\n",
    "        for feature_index in feature_indices:\n",
    "            # 我们获取模型对这个特征的预测。\n",
    "            start_logits = all_start_logits[feature_index]\n",
    "            end_logits = all_end_logits[feature_index]\n",
    "            # 这将允许我们将logits中的某些位置映射到原始上下文中的文本跨度。\n",
    "            offset_mapping = features[feature_index][\"offset_mapping\"]\n",
    "\n",
    "            # 更新最小空预测。\n",
    "            cls_index = features[feature_index][\"input_ids\"].index(tokenizer.cls_token_id)\n",
    "            feature_null_score = start_logits[cls_index] + end_logits[cls_index]\n",
    "            if min_null_score is None or min_null_score < feature_null_score:\n",
    "                min_null_score = feature_null_score\n",
    "\n",
    "            # 浏览所有的最佳开始和结束logits，为 `n_best_size` 个最佳选择。\n",
    "            start_indexes = np.argsort(start_logits)[-1 : -n_best_size - 1 : -1].tolist()\n",
    "            end_indexes = np.argsort(end_logits)[-1 : -n_best_size - 1 : -1].tolist()\n",
    "            for start_index in start_indexes:\n",
    "                for end_index in end_indexes:\n",
    "                    # 不考虑超出范围的答案，原因是索引超出范围或对应于输入ID的部分不在上下文中。\n",
    "                    if (\n",
    "                        start_index >= len(offset_mapping)\n",
    "                        or end_index >= len(offset_mapping)\n",
    "                        or offset_mapping[start_index] is None\n",
    "                        or offset_mapping[end_index] is None\n",
    "                    ):\n",
    "                        continue\n",
    "                    # 不考虑长度小于0或大于max_answer_length的答案。\n",
    "                    if end_index < start_index or end_index - start_index + 1 > max_answer_length:\n",
    "                        continue\n",
    "\n",
    "                    start_char = offset_mapping[start_index][0]\n",
    "                    end_char = offset_mapping[end_index][1]\n",
    "                    valid_answers.append(\n",
    "                        {\n",
    "                            \"score\": start_logits[start_index] + end_logits[end_index],\n",
    "                            \"text\": context[start_char: end_char]\n",
    "                        }\n",
    "                    )\n",
    "        \n",
    "        if len(valid_answers) > 0:\n",
    "            best_answer = sorted(valid_answers, key=lambda x: x[\"score\"], reverse=True)[0]\n",
    "        else:\n",
    "            # 在极少数情况下我们没有一个非空预测，我们创建一个假预测以避免失败。\n",
    "            best_answer = {\"text\": \"\", \"score\": 0.0}\n",
    "        \n",
    "        # 选择我们的最终答案：最佳答案或空答案（仅适用于squad_v2）\n",
    "        if not squad_v2:\n",
    "            predictions[example[\"id\"]] = best_answer[\"text\"]\n",
    "        else:\n",
    "            answer = best_answer[\"text\"] if best_answer[\"score\"] > min_null_score else \"\"\n",
    "            predictions[example[\"id\"]] = answer\n",
    "\n",
    "    return predictions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "正在后处理 11873 个示例的预测，这些预测分散在 12134 个特征中。\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c830b5bf7f734cd2981c9f35e695aa43",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/11873 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "final_predictions = postprocess_qa_predictions(datasets[\"validation\"], validation_features, raw_predictions.predictions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/lib/python3.10/site-packages/datasets/load.py:753: FutureWarning: The repository for squad_v2 contains custom code which must be executed to correctly load the metric. You can inspect the repository content at https://raw.githubusercontent.com/huggingface/datasets/2.17.0/metrics/squad_v2/squad_v2.py\n",
      "You can avoid this message in future by passing the argument `trust_remote_code=True`.\n",
      "Passing `trust_remote_code=True` will be mandatory to load this metric from the next major release of `datasets`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e14179f673b14a42adc5c45ab3b7f760",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading extra modules:   0%|          | 0.00/3.19k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "{'exact': 61.18083045565569,\n",
       " 'f1': 64.7949234954046,\n",
       " 'total': 11873,\n",
       " 'HasAns_exact': 64.65924426450742,\n",
       " 'HasAns_f1': 71.89779464590703,\n",
       " 'HasAns_total': 5928,\n",
       " 'NoAns_exact': 57.712363330529854,\n",
       " 'NoAns_f1': 57.712363330529854,\n",
       " 'NoAns_total': 5945,\n",
       " 'best_exact': 61.18083045565569,\n",
       " 'best_exact_thresh': 0.0,\n",
       " 'best_f1': 64.7949234954048,\n",
       " 'best_f1_thresh': 0.0}"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_metric\n",
    "\n",
    "metric = load_metric(\"squad_v2\" if squad_v2 else \"squad\")\n",
    "\n",
    "if squad_v2:\n",
    "    formatted_predictions = [{\"id\": k, \"prediction_text\": v, \"no_answer_probability\": 0.0} for k, v in final_predictions.items()]\n",
    "else:\n",
    "    formatted_predictions = [{\"id\": k, \"prediction_text\": v} for k, v in final_predictions.items()]\n",
    "references = [{\"id\": ex[\"id\"], \"answers\": ex[\"answers\"]} for ex in datasets[\"validation\"]]\n",
    "metric.compute(predictions=formatted_predictions, references=references)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 原始分数，{'exact_match': 74.88174077578051, 'f1': 83.6359321422016}\n",
    "\n",
    "全量训练后，发现f1降低，train loss在降低，但是验证集的loss在变大，为啥呢？需要分析原因并再次训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "name": "Question Answering on SQUAD",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
